#include "Multigrid.h"
#include "json.hpp"

#define exp 2.718281828

int globalInd = 0;
int rInd = 0;
int sInd = 0;
int errInd = 0;
int timeInd = 0;

// 用于存放范数误差
double errorNorm[4];

// CPU 时间
double lpTime[4];
double muTime[4];

// 绘图函数
template <int dim>
void plot(Grid<dim> &gData, Container2D &U, Container2D &V, Container2D &EX, int times, Real *residual, double t1, double t2)
{
    // 绘制解

    // 横纵坐标
    cout << "X = [\n";
    gData.print();
    cout << "];" << endl;

    // 解向量
    cout << "V = [\n"
         << V << "];" << endl;

    cout << "figure(1);" << endl;
    cout << "plot3(X(:,1),X(:,2),V, '.');" << endl;
    cout << "hold on;" << endl;

    // 绘制残量范数
    cout << "R = [\n";
    for (int i = 0; i < times; i++)
    {
        cout << log2(residual[i]) << endl;
    }
    cout << "];" << endl;

    cout << "figure(2);" << endl;
    cout << "subplot(4,1," << globalInd + 1 << ");" << endl;
    cout << "plot(R,'.');" << endl;
    cout << "title('Residual: n = " << gData.n << "');" << endl;

    // 集齐 4 张绘制
    if (globalInd + 1 == 4)
    {
        cout << "figure(1);" << endl;
        cout << "legend('4', '8', '16', '32');" << endl;
        cout << "saveas(gcf,'oneSolutionInd" << ++sInd << ".png');" << endl;
        cout << "clf" << endl;

        cout << "figure(2);" << endl;
        cout << "saveas(gcf,'oneResidualInd" << ++rInd << ".png');" << endl;
    }

    // 计算误差并汇报信息
    errorNorm[globalInd] = gData.err(V - EX, 0);

    // 时间
    lpTime[globalInd] = t1;
    muTime[globalInd] = t2;

    globalInd = (globalInd + 1) % 4;
}

// 计算函数
template <int dim>
void calculate(json &inputs, MultiSolver<dim> &ms)
{
    // 读取数据的总量
    int num = inputs["Num"];

    // 逐个计算
    for (int i = 0; i < num; i++)
    {
        // 当前数据
        json nj = inputs["Data"][i];

        int n = 4;
        int p = 4;

        for (int j = 0; j < p; j++)
        {
            // 记录规模
            ms.gData.n = n;

            // 使用无穷范数
            int norm = 0;

            // 读取边界数据
            ms.read(nj);
            ms.multiSolve(norm, plot);

            cout << endl;

            n *= 2;
        }

        // 作为横坐标
        int N = 4;
        cout << "N = [\n";
        for (int i = 0; i < p; i++)
        {
            cout << N << endl;
            N *= 2;
        }
        cout << "];" << endl;

        // 绘制残量范数
        cout << "E = [\n";
        for (int i = 0; i < p; i++)
        {
            cout << log2(errorNorm[i]) << endl;
        }
        cout << "];" << endl;

        cout << "figure(3);" << endl;
        cout << "plot(log2(N),E,'-', 'MarkerSize', 5);" << endl;
        cout << "hold on;" << endl;
        cout << endl;

        // 绘制 CPU 时间
        cout << "T1 = [\n";
        for (int i = 0; i < p; i++)
        {
            cout << lpTime[i] << endl;
        }
        cout << "];" << endl;

        cout << "T2 = [\n";
        for (int i = 0; i < p; i++)
        {
            cout << muTime[i] << endl;
        }
        cout << "];" << endl;

        cout << "figure(4);" << endl;
        cout << "subplot(3,1," << i + 1 << ");" << endl;
        cout << "plot(N,T1, '-', 'MarkerSize', 5);" << endl;
        cout << "hold on;" << endl;
        cout << "plot(N,T2, '-', 'MarkerSize', 5);" << endl;
        cout << "legend('Lapcak', 'Multigrid');" << endl;
        cout << "ylabel('time(s)');" << endl;
        cout << "xlabel('n');" << endl;
        cout << endl;
    }

    // 绘制误差收敛图
    cout << "figure(3);" << endl;
    cout << "title('Convergence rates');" << endl;
    cout << "legend('Dirichlet', 'Neumann', 'Mixed');" << endl;
    cout << "xlabel('log n');" << endl;
    cout << "ylabel('log of error');" << endl;
    cout << "saveas(gcf,'oneErrorInd" << ++errInd << ".png');" << endl;

    cout << "clf;" << endl;

    // 绘制 CPU 时间图
    cout << "figure(4);" << endl;
    cout << "subplot(3,1,1);" << endl;
    cout << "title('Dirichlet');" << endl;
    cout << "subplot(3,1,2);" << endl;
    cout << "title('Neumann');" << endl;
    cout << "subplot(3,1,3);" << endl;
    cout << "title('Mixed');" << endl;
    cout << "saveas(gcf,'oneTimeInd" << ++timeInd << ".png');" << endl;

    cout << "clf;" << endl;
}

void testOne(json &inputs)
{
    Func<2> u = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return pow(exp, y + sin(x));
    };

    Func<2> ux = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return pow(exp, y + sin(x)) * cos(x);
    };

    Func<2> uy = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return pow(exp, y + sin(x));
    };

    Func<2> f = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return -pow(exp, y + sin(x)) * (1 + cos(x) * cos(x) - sin(x));
    };

    MultiSolver<2> ms;

    // 函数
    ms.gData.f = f;
    ms.gData.u = u;
    ms.gData.du[0] = ux;
    ms.gData.du[1] = uy;

    calculate(inputs, ms);
}

void testTwo(json &inputs)
{
    Func<2> u = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return sin(x + y) - x * y;
    };

    Func<2> ux = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return cos(x + y) - y;
    };

    Func<2> uy = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return cos(x + y) - x;
    };

    Func<2> f = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return 2 * sin(x + y);
    };

    MultiSolver<2> ms;

    // 函数
    ms.gData.f = f;
    ms.gData.u = u;
    ms.gData.du[0] = ux;
    ms.gData.du[1] = uy;

    calculate(inputs, ms);
}

void testThree(json &inputs)
{
    Func<2> u = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return sin(x) + sin(y) + x * y;
    };

    Func<2> ux = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return cos(x) + y;
    };

    Func<2> uy = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return cos(y) + x;
    };

    Func<2> f = [](Point<2> p)
    {
        Real x = p[0];
        Real y = p[1];
        return sin(x) + sin(y);
    };

    MultiSolver<2> ms;

    // 函数
    ms.gData.f = f;
    ms.gData.u = u;
    ms.gData.du[0] = ux;
    ms.gData.du[1] = uy;

    calculate(inputs, ms);
}

int main()
{
    fstream fp("draw.json", ios::in);
    json inputs;
    fp >> inputs;

    cout << endl;
    testOne(inputs);

    cout << endl;
    testTwo(inputs);

    cout << endl;
    testThree(inputs);

    return 0;
}